import xml.dom.minidom
import os
import cv2

PATH_annotations = os.path.abspath('./train_annotations')  # 当前文件路径

PATH_images = os.path.abspath('./train_images')
PATH_images_to = os.path.abspath('./_train_images')

PATH_without_mask = os.path.abspath('./train/without_mask')
PATH_with_mask = os.path.abspath('./train/with_mask')
PATH_mask_weared_incorrect = os.path.abspath('./train/mask_weared_incorrect')
for i in range(200, 853):
    name = str(i)
    _list_x = []
    _list_y = []
    mask = []
    os.chdir(PATH_annotations)
    dom = xml.dom.minidom.parse(name + '.xml')
    rootdata = dom.documentElement
    item = rootdata.getElementsByTagName("object")
    image_path = (name + '.png')
    print(image_path)
    os.chdir(PATH_images)
    image = cv2.imread(image_path)
    os.chdir(PATH_images_to)
    image_to = cv2.imread(image_path)
    for num in item:  # 一个照片中的人头数
        xmin = num.getElementsByTagName("xmin")[0]
        ymin = num.getElementsByTagName("ymin")[0]
        xmax = num.getElementsByTagName("xmax")[0]
        ymax = num.getElementsByTagName("ymax")[0]
        linshi_x = [xmin.childNodes[0].data, xmax.childNodes[0].data]
        linshi_y = [ymin.childNodes[0].data, ymax.childNodes[0].data]
        lable = num.getElementsByTagName("name")[0]

        linshi_mask = lable.childNodes[0].data

        if str(linshi_mask)=='with_mask':
            # mask[num]=1
            mask.append(1)
        if str(linshi_mask)=='without_mask':
            # mask[num]=0
            mask.append(0)
        if str(linshi_mask)=='mask_weared_incorrect':
            mask.append(2)
        _list_x.append(linshi_x)
        _list_y.append(linshi_y)

    # print(len(_list_x))
    # print(mask)
    for face in range(len(_list_x)):
        first_point = (int(_list_x[face][0]), int(_list_y[face][0]))
        last_point = (int(_list_x[face][1]), int(_list_y[face][1]))

        cropped = image[int(_list_y[face][0]):int(_list_y[face][1]), int(_list_x[face][0]):int(_list_x[face][1])]  # 裁剪坐标为[y0:y1, x0:x1]

        if mask[face]==1:

            os.chdir(PATH_images_to)
            cv2.rectangle(image_to, first_point, last_point, (0, 255, 0), 2)  # 绿色(b,g,r)
            _text='with_mask'
            font = cv2.FONT_HERSHEY_PLAIN
            text = cv2.putText(image_to,_text, (int(_list_x[face][0]),int(_list_y[face][0]) - 5), font, 0.8, (0, 255, 0), 1)
            cv2.imwrite(name + '.png', image_to)

            os.chdir(PATH_with_mask)
            cv2.imwrite(name +'_'+str(face)+ '.png', cropped)
        if mask[face]==0:

            os.chdir(PATH_images_to)
            cv2.rectangle(image_to, first_point, last_point, (0, 0, 255), 2)  # 红色(b,g,r)
            _text = 'without_mask'
            font = cv2.FONT_HERSHEY_PLAIN
            text = cv2.putText(image_to,_text, (int(_list_x[face][0]),
                                                                   int(_list_y[face][0]) - 5), font, 0.8, (0, 0, 255), 1)
            cv2.imwrite(name + '.png', image_to)

            os.chdir(PATH_without_mask)
            cv2.imwrite(name +'_'+str(face)+ '.png', cropped)
        if mask[face]==2:

            os.chdir(PATH_images_to)
            cv2.rectangle(image_to, first_point, last_point, (0, 255, 255), 2)  # 黄色
            _text = 'mask_weared_incorrect'
            font = cv2.FONT_HERSHEY_PLAIN
            text = cv2.putText(image_to,_text, (int(_list_x[face][0]),
                                                                   int(_list_y[face][0]) - 5), font, 0.8, (0, 255, 255), 1)
            cv2.imwrite(name + '.png', image_to)

            os.chdir(PATH_mask_weared_incorrect)
            cv2.imwrite(name +'_'+str(face)+ '.png', cropped)

